'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 2_study_retriever_route.py
* @Time: 2025/10/31
* @All Rights Reserve By Brtc
'''
from typing import Literal

import dotenv
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableConfig, RunnablePassthrough
from langchain_openai import ChatOpenAI
from pydantic import Field, BaseModel

dotenv.load_dotenv()
class RouteQuery(BaseModel):
    """将用户查询映射到最相关的数据"""
    datasource: Literal["python_docs", "js_docs", "golang_docs"] = Field(
        description="根据用户问题，选择那个数据源最相关，以回答他们的问题"
    )

def choose_route(result:RouteQuery) -> str:
    if "python_docs" in result.datasource.lower():
        return "chain for python"
    elif "js_docs" in result.datasource.lower():
        return "chain for javascript"
    else:
        return "chain for golang"


llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
struct_llm = llm.with_structured_output(RouteQuery)
# 3、执行相应的提问,检查映射路由
question = """ 
为什么下面的代码不工作了，请帮我检查下：
int mian(){
    int a = 10;
    int b = 20;
    print("a+b=%d\n", a+b);
}
"""
res = struct_llm.invoke(question)
#2、创建逻辑链
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个班擅长将用户问题路由到适当数据源的专家，请根据用户的问题设计的编程语言将其路由到相关数据源！"),
    ("human", "{question}")
])

router = {"question":RunnablePassthrough()}|prompt|struct_llm|choose_route

# 4、选择不同的路由
print(router.invoke(question))